/*
 * Copyright (c) 2020-2099 the original author or authors. All rights reserve.
 */
package org.tianyun.cloud.datasource.routing;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.tianyun.cloud.datasource.configuration.DataSourceFactoryProperties;
import org.tianyun.cloud.datasource.configuration.RoutingDataSourceProperties;

import javax.sql.DataSource;
import java.io.Closeable;
import java.io.IOException;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * 路由数据源，根据上下文动态选择数据源
 *
 * @auther ebert_chan
 */
public class RoutingDataSource extends AbstractRoutingDataSource {

    private static final Logger LOGGER = LoggerFactory.getLogger(RoutingDataSource.class);

    private RoutingDataSourceProperties routingDataSourceProperties;

    private Map<String, DataSourceFactory> factorys;

    private DataSourceFactory defaultFactory;

    private final AtomicBoolean isShutdown = new AtomicBoolean();

    public RoutingDataSource(List<DataSourceFactory> factorys) {
        super();
        this.setFactorys(factorys);
    }

    public RoutingDataSource(RoutingDataSourceProperties routingDataSourceProperties) throws IOException {
        super();
        this.routingDataSourceProperties = routingDataSourceProperties;
        Map<String, DataSourceFactoryProperties> nodes = routingDataSourceProperties.getNodes();
        Assert.notEmpty(nodes, "Failed to create the datasource, No configuration datasource");
        List<DataSourceFactory> factorys = new ArrayList<>(nodes.size());
        for (Map.Entry<String, DataSourceFactoryProperties> entry : nodes.entrySet()) {
            DataSourceFactoryProperties routingDataSourceFactory = entry.getValue();
            String name = entry.getKey();
            DataSourceFactory cluster = new DataSourceFactory(name, routingDataSourceFactory);
            factorys.add(cluster);
        }
        this.setFactorys(factorys);
    }

    /*
     * @see org.tianyun.cloud.datasource.AbstractRoutingDataSource#resolveDataSource( )
     */
    @Override
    public DataSource resolveDataSource() throws SQLException {
        if (isClosed()) {
            throw new SQLException("DataSource " + this + " has been closed.");
        }

        DataSourceContext context = DataSourceContext.get();
        if (context == null) {
            return defaultFactory.getMaster();
        }

        DataSourceFactory cluster = null;
        String target = context.getTarget();
        if (StringUtils.hasLength(target)) {
            cluster = factorys.get(target);
        } else {
            cluster = defaultFactory;
        }
        Assert.notNull(cluster, "DataSource '" + target + "' doesn't exist.");

        if (context.isSlave() && cluster.getSlave() != null) {
            return cluster.getSlave();
        }

        return cluster.getMaster();
    }

    /**
     * 销毁数据源，释放资源
     */
    @Override
    public void close() throws IOException {
        if (!isShutdown.getAndSet(true)) {
            return;
        }
        if (routingDataSourceProperties != null) {
            close(factorys);
        }
    }

    /**
     * 判断当前数据源是否已销毁
     *
     * @return
     * @auther ebert_chan
     */
    public boolean isClosed() {
        return isShutdown.get();
    }

    /**
     * 设置数据源，可动态刷新数据源
     *
     * @param factorys
     * @auther ebert_chan
     */
    private void setFactorys(List<DataSourceFactory> factorys) {
        Assert.notEmpty(factorys, "This collection must not be empty, it must contain at least 1 element");
        Assert.noNullElements(factorys, "Collection must contain non-null elements");

        this.factorys = factorys.stream().collect(Collectors.toMap(DataSourceFactory::getName, Function.identity()));
        for (DataSourceFactory cluster : factorys) {
            if (cluster.isPrimary()) {
                defaultFactory = cluster;
                break;
            }
        }
        if (defaultFactory == null) {
            defaultFactory = factorys.get(0);
        }
    }

    private void close(Map<String, DataSourceFactory> factorys) {
        for (Map.Entry<String, DataSourceFactory> entry : factorys.entrySet()) {
            DataSourceFactory cluster = entry.getValue();
            this.closeDataSource(cluster.getMaster());
            this.closeDataSource(cluster.getSlave());
        }
    }

    private void closeDataSource(DataSource dataSource) {
        if (dataSource != null && dataSource instanceof Closeable) {
            Closeable closeable = (Closeable) dataSource;
            try {
                closeable.close();
            } catch (IOException e) {
                LOGGER.error("DataSource close exception.", e);
            }
        }
    }

}
